#include "ntddk.h"
#include "AsmCall.h"
#include "Util.h"
#include "ia32_types.h"
#include "common.h"

#include <pshpack1.h>
struct MtrrData {
	bool enabled;        //!< Whether this entry is valid
	bool fixedMtrr;      //!< Whether this entry manages a fixed range MTRR
	UCHAR type;          //!< Memory Type (such as WB, UC)
	bool reserverd1;     //!< Padding
	ULONG reserverd2;    //!< Padding
	ULONG64 range_base;  //!< A base address of a range managed by this entry
	ULONG64 range_end;   //!< An end address of a range managed by this entry
};
#include <poppack.h>

static MtrrData g_eptp_mtrr_entries[343];
static UCHAR g_eptp_mtrr_default_type;

void EptInitializeMtrrEntries() {
	MtrrData *mtrr_entries = g_eptp_mtrr_entries;

	Ia32MtrrDefaultTypeMsr default_type;
	Asm_rdmsr((ULONG32)(Msr::kIa32MtrrDefType), &default_type.all);
	g_eptp_mtrr_default_type = default_type.fields.default_mtemory_type;

	Ia32MtrrCapabilitiesMsr mtrr_capabilities;
	Asm_rdmsr((ULONG32)(Msr::kIa32MtrrCap), &mtrr_capabilities.all);
	ULONG index = 0;

	if (mtrr_capabilities.fields.fixed_range_supported && default_type.fields.fixed_mtrrs_enabled)
	{
		ULONG64 k64kBase = 0x0;
		ULONG64 k64kManagedSize = 0x10000;
		ULONG64 k16kBase = 0x80000;
		ULONG64 k16kManagedSize = 0x4000;
		ULONG64 k4kBase = 0xC0000;
		ULONG64 k4kManagedSize = 0x1000;

		ULONG64 offset = 0;
		Ia32MtrrFixedRangeMsr fixed_range;
		Asm_rdmsr((ULONG32)(Msr::kIa32MtrrFix64k00000), &fixed_range.all);

		for (ULONG i = 0; i < (sizeof(fixed_range.fields.types) / sizeof(*(fixed_range.fields.types))); i++)
		{
			ULONG64 base = k64kBase + offset;
			offset += k64kManagedSize;

			mtrr_entries[index].enabled = true;
			mtrr_entries[index].fixedMtrr = true;
			mtrr_entries[index].type = fixed_range.fields.types[i];
			mtrr_entries[index].range_base = base;
			mtrr_entries[index].range_end = base + k64kManagedSize - 1;
			index++;
		}

		offset = 0;
		for (ULONG32 nMsr = (ULONG32)Msr::kIa32MtrrFix16k80000; \
			nMsr <= (ULONG32)Msr::kIa32MtrrFix16kA0000; nMsr++)
		{
			Asm_rdmsr(nMsr, &fixed_range.all);
			for (ULONG i = 0; i < (sizeof(fixed_range.fields.types) / sizeof(*(fixed_range.fields.types))); i++)
			{
				ULONG64 base = k16kBase + offset;
				offset += k16kManagedSize;

				mtrr_entries[index].enabled = true;
				mtrr_entries[index].fixedMtrr = true;
				mtrr_entries[index].type = fixed_range.fields.types[i];
				mtrr_entries[index].range_base = base;
				mtrr_entries[index].range_end = base + k16kManagedSize - 1;
				index++;
			}
		}

		offset = 0;
		for (ULONG32 nMsr = (ULONG32)Msr::kIa32MtrrFix4kC0000; \
			nMsr <= (ULONG32)Msr::kIa32MtrrFix4kF8000; nMsr++)
		{
			Asm_rdmsr(nMsr, &fixed_range.all);
			for (ULONG i = 0; i < (sizeof(fixed_range.fields.types) / sizeof(*(fixed_range.fields.types))); i++)
			{
				ULONG64 base = k4kBase + offset;
				offset += k4kManagedSize;

				mtrr_entries[index].enabled = true;
				mtrr_entries[index].fixedMtrr = true;
				mtrr_entries[index].type = fixed_range.fields.types[i];
				mtrr_entries[index].range_base = base;
				mtrr_entries[index].range_end = base + k4kManagedSize - 1;
				index++;
			}
		}
	}

	for (ULONG i = 0; i < mtrr_capabilities.fields.variable_range_count; i++)
	{
		ULONG32 phy_msk = (ULONG32)Msr::kIa32MtrrPhysMaskN + i * 2;
		Ia32MtrrPhysMaskMsr mtrr_mask;
		Asm_rdmsr(phy_msk, &mtrr_mask.all);
		if (!mtrr_mask.fields.valid)
		{
			continue;
		}

		ULONG length;
		BitScanForward64(&length, mtrr_mask.fields.phys_mask * PAGE_SIZE);

		ULONG32 phy_base = (ULONG32)Msr::kIa32MtrrPhysBaseN + i * 2;
		Ia32MtrrPhysBaseMsr mtrr_base;
		Asm_rdmsr(phy_base, &mtrr_base.all);
		ULONG64 base = mtrr_base.fields.phys_base * PAGE_SIZE;
		ULONG64 end = base + (1ull << length) - 1;

		mtrr_entries[index].enabled = true;
		mtrr_entries[index].fixedMtrr = false;
		mtrr_entries[index].type = mtrr_base.fields.type;
		mtrr_entries[index].range_base = base;
		mtrr_entries[index].range_end = end;
		index++;
	}
}

void EptpFreeUnusedPreAllocatedEntries(
	EptCommonEntry **preallocated_entries, long used_count) {
	for (auto i = used_count; i < 50; ++i) {
		if (!preallocated_entries[i]) {
			break;
		}
#pragma warning(push)
#pragma warning(disable : 6001)
		ExFreePoolWithTag(preallocated_entries[i], kHyperPlatformCommonPoolTag);
#pragma warning(pop)
	}
	ExFreePoolWithTag(preallocated_entries, kHyperPlatformCommonPoolTag);
}

void EptTermination(EptData *ept_data) {

	EptpFreeUnusedPreAllocatedEntries(ept_data->preallocated_entries,ept_data->preallocated_entries_count);
	EptpDestructTables(ept_data->ept_pml4, 4);
	ExFreePoolWithTag(ept_data, kHyperPlatformCommonPoolTag);
}

void EptpDestructTables(EptCommonEntry *table, ULONG table_level) {
	for (auto i = 0ul; i < 512; ++i) {
		const auto entry = table[i];
		if (entry.fields.physial_address) {
			const auto sub_table = static_cast<EptCommonEntry *>(
				UtilVaFromPfn(entry.fields.physial_address));

			switch (table_level) {
			case 4:  // table == PML4, sub_table == PDPT
			case 3:  // table == PDPT, sub_table == PDT
				EptpDestructTables(sub_table, table_level - 1);
				break;
			case 2:  // table == PDT, sub_table == PT
				ExFreePoolWithTag(sub_table, kHyperPlatformCommonPoolTag);
				break;
			default:
				break;
			}
		}
	}
	ExFreePoolWithTag(table, kHyperPlatformCommonPoolTag);
}

memory_type EptGetMemoryType(ULONG64 physical_address) {
	UCHAR result_type = MAXUCHAR;
	for (ULONG nIndex = 0; nIndex < (sizeof(g_eptp_mtrr_entries) / sizeof(*g_eptp_mtrr_entries)); nIndex++)
	{
		MtrrData mtrr_entry = g_eptp_mtrr_entries[nIndex];
		if (!mtrr_entry.enabled)
		{
			break;
		}

		if (!UtilIsInBounds(physical_address, mtrr_entry.range_base, mtrr_entry.range_end))
		{
			continue;
		}

		if (mtrr_entry.fixedMtrr)
		{
			result_type = mtrr_entry.type;
			break;
		}

		if (mtrr_entry.type == (UCHAR)(memory_type::kUncacheable))
		{
			result_type = mtrr_entry.type;
			break;
		}

		if (result_type == (UCHAR)(memory_type::kWriteThrough) || \
			mtrr_entry.type == (UCHAR)(memory_type::kWriteThrough)) {
			if (result_type == (UCHAR)(memory_type::kWriteBack))
			{
				result_type = (UCHAR)(memory_type::kWriteThrough);
				continue;
			}
		}

		result_type = mtrr_entry.type;
	}

	if (result_type == MAXUCHAR)
	{
		result_type = g_eptp_mtrr_default_type;
	}

	return (memory_type)result_type;
}

ULONG64 EptpAddressToPxeIndex(ULONG64 physical_address) {
	return (physical_address >> 39) & 0x1ff;
}

// Return an address of PPE
ULONG64 EptpAddressToPpeIndex(ULONG64 physical_address) {
	return (physical_address >> 30) & 0x1ff;
}

ULONG EptAddressToPdeIndex(ULONG64 physical_address) {
	return (physical_address >> 21) & 0x1ff;
}

ULONG EptAddressToPteIndex(ULONG64 physical_address) {
	return (physical_address >> 12) & 0x1ff;
}

EptCommonEntry *EptAllocateEptEntryFromPreAllocate(EptData *ept_data) {
	ULONG count = InterlockedIncrement(&ept_data->preallocated_entries_count);
	return ept_data->preallocated_entries[count - 1];
}

EptCommonEntry *EptAllocateEptEntryFromPool() {
	ULONG kAllocSize = 512 * sizeof(EptCommonEntry);
	EptCommonEntry *entry = (EptCommonEntry *)ExAllocatePoolWithTag(NonPagedPool, kAllocSize, kHyperPlatformCommonPoolTag);
	if (!entry)
	{
		return entry;
	}

	RtlZeroMemory(entry, kAllocSize);
	return entry;
}

EptCommonEntry *EptAllocateEptEntry(EptData *ept_data) {
	if (ept_data) {
		return EptAllocateEptEntryFromPreAllocate(ept_data);
	}
	else {
		return EptAllocateEptEntryFromPool();
	}
}

memory_type EptpGetMemoryType(ULONG64 physical_address) {
	// Indicate that MTRR is not defined (as a default)
	UCHAR result_type = MAXUCHAR;

	// Looks for MTRR that includes the specified physical_address
	for (const auto mtrr_entry : g_eptp_mtrr_entries) {
		if (!mtrr_entry.enabled) {
			// Reached out the end of stored MTRRs
			break;
		}

		if (!UtilIsInBounds(physical_address, mtrr_entry.range_base,
			mtrr_entry.range_end)) {
			// This MTRR does not describe a memory type of the physical_address
			continue;
		}

		// See: MTRR Precedences
		if (mtrr_entry.fixedMtrr) {
			// If a fixed MTRR describes a memory type, it is priority
			result_type = mtrr_entry.type;
			break;
		}

		if (mtrr_entry.type == static_cast<UCHAR>(memory_type::kUncacheable)) {
			// If a memory type is UC, it is priority. Do not continue to search as
			// UC has the highest priority
			result_type = mtrr_entry.type;
			break;
		}

		if (result_type == static_cast<UCHAR>(memory_type::kWriteThrough) ||
			mtrr_entry.type == static_cast<UCHAR>(memory_type::kWriteThrough)) {
			if (result_type == static_cast<UCHAR>(memory_type::kWriteBack)) {
				// If two or more MTRRs describes an over-wrapped memory region, and
				// one is WT and the other one is WB, use WT. However, look for other
				// MTRRs, as the other MTRR specifies the memory address as UC, which is
				// priority.
				result_type = static_cast<UCHAR>(memory_type::kWriteThrough);
				continue;
			}
		}

		// Otherwise, processor behavior is undefined. We just use the last MTRR
		// describes the memory address.
		result_type = mtrr_entry.type;
	}

	// Use the default MTRR if no MTRR entry is found
	if (result_type == MAXUCHAR) {
		result_type = g_eptp_mtrr_default_type;
	}

	return static_cast<memory_type>(result_type);
}

void EptpInitTableEntry(EptCommonEntry *entry, ULONG table_level, ULONG64 physical_address) {
	entry->fields.read_access = true;
	entry->fields.write_access = true;
	entry->fields.execute_access = true;
	entry->fields.physial_address = UtilPfnFromPa(physical_address);
	if (table_level == 1) {
		entry->fields.memory_type =
			static_cast<ULONG64>(EptpGetMemoryType(physical_address));
	}
}

EptCommonEntry *EptGetEptPtEntry(EptData *ept_data, ULONG64 physical_address) {
	return EptpGetEptPtEntry(ept_data->ept_pml4, 4, physical_address);
}

// Returns an EPT entry corresponds to the physical_address
EptCommonEntry *EptpGetEptPtEntry(EptCommonEntry *table, ULONG table_level, ULONG64 physical_address) {
	if (!table) {
		return nullptr;
	}
	switch (table_level) {
	case 4: {
		// table == PML4
		const auto pxe_index = EptpAddressToPxeIndex(physical_address);
		const auto ept_pml4_entry = &table[pxe_index];
		if (!ept_pml4_entry->all) {
			return nullptr;
		}
		return EptpGetEptPtEntry(static_cast<EptCommonEntry *>(UtilVaFromPfn(
			ept_pml4_entry->fields.physial_address)),
			table_level - 1, physical_address);
	}
	case 3: {
		// table == PDPT
		const auto ppe_index = EptpAddressToPpeIndex(physical_address);
		const auto ept_pdpt_entry = &table[ppe_index];
		if (!ept_pdpt_entry->all) {
			return nullptr;
		}
		return EptpGetEptPtEntry(static_cast<EptCommonEntry *>(UtilVaFromPfn(
			ept_pdpt_entry->fields.physial_address)),
			table_level - 1, physical_address);
	}
	case 2: {
		// table == PDT
		const auto pde_index = EptAddressToPdeIndex(physical_address);
		const auto ept_pdt_entry = &table[pde_index];
		if (!ept_pdt_entry->all) {
			return nullptr;
		}
		return EptpGetEptPtEntry(static_cast<EptCommonEntry *>(UtilVaFromPfn(
			ept_pdt_entry->fields.physial_address)),
			table_level - 1, physical_address);
	}
	case 1: {
		// table == PT
		const auto pte_index = EptAddressToPteIndex(physical_address);
		const auto ept_pt_entry = &table[pte_index];
		return ept_pt_entry;
	}
	default:
		//HYPERPLATFORM_COMMON_DBG_BREAK();
		return nullptr;
	}
}

EptCommonEntry *EptpConstructTables(EptCommonEntry *table, ULONG table_level, ULONG64 physical_address,EptData *ept_data) {
	if (table == 0)
	{
		return 0;
	}
	if (table_level == 4)
	{
		ULONG pxe_index = (ULONG)EptpAddressToPxeIndex(physical_address);
		EptCommonEntry *ept_pml4_entry = &table[pxe_index];
		if (!ept_pml4_entry->all)
		{
			EptCommonEntry *ept_pdpt = EptAllocateEptEntry(ept_data);
			if (!ept_pdpt)
			{
				return 0;
			}

			EptpInitTableEntry(ept_pml4_entry, table_level, UtilPaFromVa(ept_pdpt));
		}
		return EptpConstructTables((EptCommonEntry *)UtilVaFromPfn(ept_pml4_entry->fields.physial_address), table_level - 1, physical_address, ept_data);
	}
	else if (table_level == 3) {
		ULONG ppe_index = (ULONG)EptpAddressToPpeIndex(physical_address);
		EptCommonEntry *ept_pdpt_entry = &table[ppe_index];
		if (!ept_pdpt_entry->all)
		{
			EptCommonEntry *ept_pdt = EptAllocateEptEntry(ept_data);
			if (!ept_pdt)
			{
				return 0;
			}

			EptpInitTableEntry(ept_pdpt_entry, table_level, UtilPaFromVa(ept_pdt));
		}

		return EptpConstructTables((EptCommonEntry *)UtilVaFromPfn(ept_pdpt_entry->fields.physial_address), table_level - 1, physical_address, ept_data);
	}
	else if (table_level == 2) {
		ULONG pde_index = EptAddressToPdeIndex(physical_address);
		EptCommonEntry *ept_pdt_entry = &table[pde_index];
		if (!ept_pdt_entry->all)
		{
			EptCommonEntry *ept_pt = EptAllocateEptEntry(ept_data);
			if (!ept_pt)
			{
				return 0;
			}
			EptpInitTableEntry(ept_pdt_entry, table_level, UtilPaFromVa(ept_pt));
		}
		return EptpConstructTables((EptCommonEntry *)UtilVaFromPfn(ept_pdt_entry->fields.physial_address), table_level - 1, physical_address, ept_data);
	}
	else if (table_level == 1) {
		ULONG pte_index = EptAddressToPteIndex(physical_address);
		EptCommonEntry *ept_pt_entry = &table[pte_index];
		EptpInitTableEntry(ept_pt_entry, table_level, physical_address);
		return ept_pt_entry;
	}
	else {
		return 0;
	}
}

void EptpFreeUnusedPreAllocateEntries(EptCommonEntry **preallocated_entries, ULONG used_count) {
	for (ULONG i = used_count; i < 50; i++)
	{
		if (!preallocated_entries[i])
		{
			break;
		}

		ExFreePoolWithTag(preallocated_entries[i], kHyperPlatformCommonPoolTag);
	}

	ExFreePoolWithTag(preallocated_entries, kHyperPlatformCommonPoolTag);
}

EptData *EptInitialization() {


	ULONG kEptPageWalkLevel = 4;
	EptData *ept_data = (EptData *)ExAllocatePoolWithTag(NonPagedPool, sizeof(EptData), kHyperPlatformCommonPoolTag);
	if (!ept_data)
	{
		KdPrint(("[Whisper]Allocate ept_data error!\n"));
		return 0;
	}

	RtlZeroMemory(ept_data, sizeof(EptData));

	EptCommonEntry *ept_pml4 = (EptCommonEntry *)ExAllocatePoolWithTag(NonPagedPool, PAGE_SIZE, kHyperPlatformCommonPoolTag);
	if (!ept_pml4)
	{
		KdPrint(("[Whisper]Allocate ept_pml4 error!\n"));
		ExFreePoolWithTag(ept_data, kHyperPlatformCommonPoolTag);
		return 0;
	}
	RtlZeroMemory(ept_pml4, PAGE_SIZE);
	ept_data->ept_pointer.all = 0;
	ept_data->ept_pointer.fields.memory_type = (ULONG64)EptGetMemoryType(UtilPaFromVa(ept_pml4));
	ept_data->ept_pointer.fields.page_walk_length = kEptPageWalkLevel - 1;
	ept_data->ept_pointer.fields.pml4_address = UtilPfnFromPa(UtilPaFromVa(ept_pml4));

	PhysicalMemoryDescriptor *pm_ranges = (PhysicalMemoryDescriptor *)UtilGetPhysicalMemoryRanges();
	for (ULONG64 run_index = 0; run_index < pm_ranges->number_of_runs; ++run_index)
	{
		PhysicalMemoryRun *run = &pm_ranges->run[run_index];
		ULONG_PTR base_addr = run->base_page * PAGE_SIZE;
		for (ULONG64 page_index = 0; page_index < run->page_count; page_index++)
		{
			ULONG_PTR indexed_addr = base_addr + page_index * PAGE_SIZE;
			EptCommonEntry *ept_pt_entry = EptpConstructTables(ept_pml4, 4, indexed_addr, 0);
			if (!ept_pt_entry)
			{
				KdPrint(("[Whisper]ConstructTable ept_pt_entry error!\n"));
				EptpDestructTables(ept_pml4, 4);
				ExFreePoolWithTag(ept_data, kHyperPlatformCommonPoolTag);
				return 0;
			}
		}
	}

	Ia32ApicBaseMsr apic_msr;
	Asm_rdmsr((ULONG32)Msr::kIa32ApicBase, &apic_msr.all);
	if (!EptpConstructTables(ept_pml4, 4, apic_msr.fields.apic_base * PAGE_SIZE, 0))
	{
		KdPrint(("[Whisper]ConstructTable apic_base error!\n"));
		EptpDestructTables(ept_pml4, 4);
		ExFreePoolWithTag(ept_data, kHyperPlatformCommonPoolTag);
		return 0;
	}

	ULONG32 preallocated_entries_size = sizeof(EptCommonEntry *) * 50;
	EptCommonEntry **preallocate_entries = (EptCommonEntry **)ExAllocatePoolWithTag(NonPagedPool, preallocated_entries_size, kHyperPlatformCommonPoolTag);
	if (!preallocate_entries)
	{
		KdPrint(("[Whisper]Allocate preallocate_entries error!\n"));
		EptpDestructTables(ept_pml4, 4);
		ExFreePoolWithTag(ept_data, kHyperPlatformCommonPoolTag);
		return 0;
	}

	RtlZeroMemory(preallocate_entries, preallocated_entries_size);

	for (ULONG i = 0; i < 50; i++)
	{
		EptCommonEntry *ept_entry = EptAllocateEptEntry(0);
		if (!ept_entry)
		{
			KdPrint(("[Whisper]Allocate ept_entry error!\n"));
			EptpFreeUnusedPreAllocateEntries(preallocate_entries, 0);
			EptpDestructTables(ept_pml4, 4);
			ExFreePoolWithTag(ept_data, kHyperPlatformCommonPoolTag);
			return 0;
		}
		preallocate_entries[i] = ept_entry;
	}

	ept_data->ept_pml4 = ept_pml4;
	ept_data->preallocated_entries = preallocate_entries;
	ept_data->preallocated_entries_count = 0;
	return ept_data;
}

void EnableMTF()
{
	ULONG64 uCPUBase;
	uCPUBase = Asm_VmRead((ULONG_PTR)VmcsField::kCpuBasedVmExecControl);
	((VmxProcessorBasedControls*)(&uCPUBase))->fields.monitor_trap_flag = 1;
	Asm_vmxWrite((ULONG_PTR)VmcsField::kCpuBasedVmExecControl, uCPUBase);
}

void DisableMTF()
{
	ULONG64 uCPUBase;
	uCPUBase = Asm_VmRead((ULONG_PTR)VmcsField::kCpuBasedVmExecControl);
	((VmxProcessorBasedControls*)(&uCPUBase))->fields.monitor_trap_flag = 0;
	Asm_vmxWrite((ULONG_PTR)VmcsField::kCpuBasedVmExecControl, uCPUBase);
}